package com.blog4j.requestRate.aop;

import com.blog4j.requestRate.annotation.RequestRate;
import com.blog4j.requestRate.exception.RequestRateException;
import com.blog4j.requestRate.properties.RequestRateProperties;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.lang.reflect.Method;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * @author 98k灬
 * @version v1.0.0
 * @Description : 功能描述
 * @Create on : 2024/9/1 11:19
 **/
@Slf4j
@Aspect
public class RequestRateAspect {
    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    RequestRateProperties properties;

    @Autowired
    private void setProperties(RequestRateProperties properties) {
        this.properties = properties;
    }

    @Pointcut("@annotation(com.blog4j.requestRate.annotation.RequestRate)")
    public void rate() {}

    @Before("rate()")
    public void doBefore(JoinPoint joinPoint) {
        if (!properties.getEnabled()) {
            return;
        }

        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        RequestRate annotation = method.getAnnotation(RequestRate.class);
        if (Objects.isNull(annotation)) {
            return;
        }

        String keyPrefix = properties.getKeyPrefix();
        String requestRateKey = StringUtils.isNotBlank(keyPrefix) ? keyPrefix : annotation.requestRateKey();

        TimeUnit mappingTimeUnit = mappingTimeUnit(properties.getTimeUnit());
        TimeUnit timeUnit = Objects.nonNull(mappingTimeUnit) ? mappingTimeUnit : annotation.timeUnit();
        int count = annotation.count();
        int time = annotation.time();
        String errMsg = annotation.errMsg();
        String className = method.getDeclaringClass().getName();

        String cacheKey = requestRateKey + "_" + className + "_" + method.getName();
        String cacheVal = stringRedisTemplate.opsForValue().get(cacheKey);
        if (StringUtils.isBlank(cacheVal)) {
            stringRedisTemplate.opsForValue().set(cacheKey, "1", time, timeUnit);
            return;
        }

        int val = Integer.parseInt(cacheVal);
        if (val >= count) {
            throw new RequestRateException(errMsg);
        }
        stringRedisTemplate.opsForValue().increment(cacheKey);
    }

    private TimeUnit mappingTimeUnit(String propertiesTimeUnit) {
        if (StringUtils.isBlank(propertiesTimeUnit)) {
            return null;
        }

        TimeUnit timeUnit;
        if (StringUtils.equalsIgnoreCase(propertiesTimeUnit, "MILLISECONDS")) {
            timeUnit = TimeUnit.MILLISECONDS;
        } else if (StringUtils.equalsIgnoreCase(propertiesTimeUnit, "MINUTE")) {
            timeUnit = TimeUnit.MINUTES;
        } else if (StringUtils.equalsIgnoreCase(propertiesTimeUnit, "HOUR")) {
            timeUnit = TimeUnit.HOURS;
        } else {
            timeUnit = TimeUnit.SECONDS;
        }
        return timeUnit;
    }
}
